import math
import time
from itertools import product
import numpy as np
from sklearn.neighbors._kde import KernelDensity
import matplotlib.pyplot as plt
import seaborn as sea
import pandas as pd

def cal_joint_prob(X, Y, O_i, parents, args):
    """
    计算P(Oi=X,pa(Oi)=Y)的值
    :param O_i:变量Oi的值 应该是:L*H 维度或者 L*H*1维度
    :param parents:变量Oi父母的值 L*H*父母个数维度
    :param args:参数
    :return:返回对应的概率值
    """

    result = 0
    for t in range(args.h):
        for l in range(args.L):
            # 这里对应exp(-w||Oi-o||2)
            tmp1 = np.exp(-args.w * np.linalg.norm(X - O_i[l, t]))
            tmp_total = 0
            for parent_i, y in enumerate(Y):
                tmp_total += np.linalg.norm(y - parents[l, t, parent_i])
            tmp2 = np.exp(-args.w * tmp_total)
            # result 等于 exp(-w||Oi -oi||2 - w\sum ||X-x||2)
            result += (tmp1 * tmp2)
    result = ((args.alpha * args.w) / (args.h * args.L)) * result
    return result


def cal_marginal_pro(parents, parents_values, args):
    """
    计算P(pa(Oi)=Y)的概率值
    :param parents: Y值
    :param parents_values: pa(Oi)中每个parent变量的值
    :param args:参数
    :return:返回对应的概率值
    """
    result = 0
    for t in range(args.h):
        for l in range(args.L):
            tmp_total = 0
            for parent_i, y in enumerate(parents):
                tmp_total += np.linalg.norm(y - parents_values[l, t, parent_i])
            tmp2 = np.exp(-args.w * tmp_total)
            # result 等于 exp(- w\sum ||X-x||2)
            result += tmp2
    result = ((args.alpha * args.w) / (args.h * args.L)) * result
    return result


def cal_condition_entropy(marginal_prob: cal_marginal_pro, joint_prob: cal_joint_prob, all_variable_set, args):
    """
    计算条件熵 H(O^i|pa(O^i))
    :param marginal_prob:计算边缘概率分布的函数
    :param joint_prob:计算联合分布的函数
    :param all_variable_set: 所有变量的取值矩阵 维度：L*H*N L是轨迹数，H是每个轨迹的时间步,N为变量个数
    :return: 返回H(X|Y) == H(O^i|pa(O^i))
    """
    entropy = 0
    for trajectory in all_variable_set:
        kde = KernelDensity(kernel='gaussian', bandwidth=0.5).fit(trajectory)
        for value in product(*trajectory.T):
            # 这里的value就是(x, y1 ,y2,...,ym)
            all_values = np.array(value, copy=True)
            # 将x取出
            x = np.delete(all_values, args.i)
            # 计算P(x,y1,y2,...,ym)
            joint = np.exp(kde.score(np.array(value).reshape(1, -1)))
            # joint = joint_prob(x, all_values, all_variable_set[:, :, args.i],all_variable_set, args)
            # 计算P(y1,y2,...,ym)
            # marginal = marginal_prob(all_values, all_variable_set, args)
            marginal = np.exp(kde.score(all_values.reshape(1, -1)))
            entropy += joint * np.log2(joint / marginal)
    print("entropy:{}".format(entropy))
    return entropy




if __name__ == '__main__':
    # zz = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9],  [10, 11, 0]])
    # print(zz.shape)
    # for item in product(*zz.T):
    #     print(item)
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument('--h', type=int, default=100, help="每一条轨迹的时间步数")
    parser.add_argument('--L', type=int, default=100, help="多少条轨迹")
    parser.add_argument('--alpha', type=int, default=1, help="normalization")
    parser.add_argument('--w', type=int, default=0.3, help="带宽")
    args = parser.parse_args()

    obs = np.random.normal(size=(10, 100, 4))
    args.L = 10
    args.h = 100
    entropy = []
    start_time = time.time()
    for i in range(4):
        args.i = i
        data = np.concatenate((obs[:, :, i].reshape((args.L, args.h, 1)), obs), axis=2)
        print(data.shape)
        entropy.append(cal_condition_entropy(cal_marginal_pro, cal_joint_prob, data, args))
    print("total time:{}".format(time.time() - start_time))
    print(entropy)
